import torch
from transformers import AutoModelForCausalLM

from deepseek_vl.models import VLChatProcessor, MultiModalityCausalLM
from deepseek_vl.utils.io import load_pil_images

#需要设置 HF_HUB_CACHE HF_HOME
#https://huggingface.co/deepseek-ai/deepseek-vl-1.3b-base/tree/main
#python3 -m huggingface_transformers.test -W ./huggingface_transformers
if __name__ == "__main__":
    # specify the path to the model
    model_path = "deepseek-ai/deepseek-vl-1.3b-base"
    vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
    tokenizer = vl_chat_processor.tokenizer

    vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
    vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()

    ## single image conversation example
    conversation = [
        {
            "role": "User",
            "content": "<image_placeholder>描述一下图中的每个步骤",
            "images": ["../DeepSeek-VL/images/training_pipelines.jpg"],
        },
        {"role": "Assistant", "content": ""},
    ]

    ## multiple images (or in-context learning) conversation example
    # conversation = [
    #     {
    #         "role": "User",
    #         "content": "<image_placeholder>A dog wearing nothing in the foreground, "
    #                    "<image_placeholder>a dog wearing a santa hat, "
    #                    "<image_placeholder>a dog wearing a wizard outfit, and "
    #                    "<image_placeholder>what's the dog wearing?",
    #         "images": [
    #             "../DeepSeek-VL/images/dog_a.png",
    #             "../DeepSeek-VL/images/dog_b.png",
    #             "../DeepSeek-VL/images/dog_c.png",
    #             "../DeepSeek-VL/images/dog_d.png",
    #         ],
    #     },
    #     {"role": "Assistant", "content": ""}
    # ]

    # load images and prepare for inputs
    pil_images = load_pil_images(conversation)
    prepare_inputs = vl_chat_processor(
        conversations=conversation,
        images=pil_images,
        force_batchify=True
    ).to(vl_gpt.device)

    # run image encoder to get the image embeddings
    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)

    # run the model to get the response
    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=512,
        do_sample=False,
        use_cache=True
    )

    for output in outputs:
        answer = tokenizer.decode(output.cuda().tolist(), skip_special_tokens=True)
        for prepare_input in prepare_inputs['sft_format']:
            print(f"{prepare_input}", answer)